#include <iostream>
#include <nlohmann/json.hpp>
#include <smart_serve_client.h>
#include <string>

using json = nlohmann::json;

using namespace OHOS::SmartServe;

template<typename Object>
std::string Dump(const Object& obj)
{
    return json(obj).dump();
}

auto MakeEvalLLMRequest(const std::string& model, const std::string& input)
{
    json request = {
        {"model", model},
        {"input", input},
    };
    return request.dump();
}

auto MakeGetLLMResultRequest(uint64_t id)
{
    json request = {
        {"id", id},
    };
    return request.dump();
}

static std::string GetResponse(const std::string& response, std::string& error)
{
    std::string content = "";
    try {
        json j = json::parse(response);
        if (j.contains("content")) {
            content = j["content"];
            error = j["error"];
        } else {
            error = "Aborted";
        }
    } catch (const json::parse_error& e) {
        error = "Aborted";
    }
    return content;
}

static uint64_t GetTaskid(const std::string& result)
{
    uint64_t id = 0;
    try {
        json data = json::parse(result);
        if (data.contains("id")) {
            id = data["id"];
        } else {
            id = 0;
        }
    } catch (const json::parse_error& e) {
        id = 0;
    }
    return id;
}

int main()
{
    std::string model = "qwen3-0.6b";
    auto& client = SmartServeClient::GetInstance();
    std::cout << "Chatbox started. Type 'exit' to quit.\n";
    std::string input;
    while (true) {
        std::cout << "> ";
        std::getline(std::cin, input);
        if (input == "exit") {
            break;
        }
        std::string text = "<|im_start|>user\n" + input + "<|im_end|>\n<|im_start|>assistant\n";
        std::string result = client.Serve(MakeEvalLLMRequest(model, text));
        uint64_t id = GetTaskid(result);
        if (id == 0) {
            std::cout << "Failed to send request by Serve interface" << std::endl;
            break;
        }
        std::string error = "";
        std::string response = "";
        std::cout << "AI: " << std::flush;
        while (error == "") {
            result = client.Serve(MakeGetLLMResultRequest(id));
            response = GetResponse(result, error);
            if (error == "Aborted") {
                std::cout << "Inference error" << std::endl;
                break;
            }
            response = GetResponse(result, error);
            std::cout << response << std::flush;
        }
        if (error == "Finished") {
            std::cout << std::endl;
        }
    }
}